import gymnasium
import numpy as np
from gymnasium import spaces
from trade_account import TradeAccount  # 你自己的账户类
from sklearn.preprocessing import StandardScaler

class TradeEnv(gymnasium.Env):
    """
    加密货币交易环境，状态是过去 seq_len 个时间步的行情 + 账户信息序列。
    每个训练回合（episode）用滑动窗口顺序覆盖全数据，避免只训练数据开头。
    """

    def __init__(
            self,
            seq_len,
            df,
            tech_indicator_list=None,
            account=None,
    ):
        super().__init__()

        self.seq_len = seq_len
        self.account = account if account is not None else TradeAccount()

        self.df_fields = ['open', 'high', 'low', 'close', 'volume'] + (tech_indicator_list or [])

        self.original_data = df[self.df_fields].values.astype(np.float32)
        self.observation_data = StandardScaler().fit_transform(self.original_data)

        self.action_space = spaces.Discrete(5)


        self.observation_space = spaces.Dict({
            "features": spaces.Box(
                low=-np.inf,
                high=np.inf,
                shape=(
                    self.seq_len,
                    len(self.df_fields) +
                    len(self.account.get_account_state())
                ),
                dtype=np.float32
            ),
            "action_mask": spaces.MultiBinary(self.action_space.n) # type: ignore
        })

        self.data_index = self.seq_len - 1 # 初始指针位置

        self.bout_total_reward = 0 # 当前回合总分数
        self.bout_total_steps = 0 # 当前回合总步数
        self.max_bout_total_steps = 0 # 当前回合总步数
        self.bout_invalid_actions = 0 # 当前回合无效操作次数

        self.action_count = {0: 0, 1: 0, 2: 0, 3: 0, 4: 0}  # 初始化操作计数

        self.account_state_history = {} # 账户状态历史 - 键为索引(data_index)，值为账户状态

        self.reward_stats = {} # 奖惩行为记录系统 - 记录分数
        self.bout_stats = {}  # 奖惩事件分数记录


    def reset(self, *, seed=None, options=None, initial_data=None):
        super().reset(seed=seed)
        self.account.reset()
        self.bout_total_reward = 0
        self.max_bout_total_steps = max(self.max_bout_total_steps, self.bout_total_steps)  # 更新最大步数
        self.bout_total_steps = 0
        self.bout_invalid_actions = 0
        self.account_state_history = {}
        self.action_count = {0: 0, 1: 0, 2: 0, 3: 0, 4: 0}  # 重置操作计数

        # 重置奖惩统计
        for key in self.reward_stats:
            self.reward_stats[key] = 0.0

        if self.data_index >= len(self.observation_data) - self.seq_len - 1:
            self.data_index = self.seq_len - 1
        else:
            self.data_index += self.seq_len

        return self.get_obs(), {}

    def get_obs(self,action_mask = None):
        if action_mask is None:
            action_mask = self.get_action_mask()
        action_mask = action_mask.astype(np.bool_)
        return {
            "features": self._get_observation(),
            "action_mask": action_mask
        }

    def _get_observation(self):
        start_idx = self.data_index - self.seq_len + 1
        end_idx = self.data_index + 1
        obs_seq = []

        for idx in range(start_idx, end_idx):
            current_data = self.observation_data[idx]
            account_state = self.account_state_history.get(idx, self.account.get_account_state())
            obs_seq.append(np.concatenate([current_data, account_state]))

        obs = np.array(obs_seq, dtype=np.float32)
        if np.any(np.isnan(obs)) or np.any(np.isinf(obs)):
            raise ValueError("观测数据包含NaN或无限值")
        return obs

    def get_action_mask(self):
        return np.array([
            True,  # 动作0: 无操作
            self.account.is_can_open_long(),
            self.account.is_can_open_short(),
            self.account.is_can_close_long(),
            self.account.is_can_close_short(),
        ], dtype=bool)

    def _record_reward_event(self, event_type, reward_value):
        """记录奖惩事件分数"""
        if event_type not in self.reward_stats:
            self.reward_stats[event_type] = 0.0

        self.reward_stats[event_type] += reward_value
    def _record_bout_event(self, event_type):
        """记录奖惩事件分数"""
        if event_type not in self.bout_stats:
            self.bout_stats[event_type] = 0.0
        self.bout_stats[event_type] += 1

    def get_volatility(self, volatility_lookback = 60, volatility_threshold = 0.002): # 判断波动率
        start_idx = max(self.data_index - volatility_lookback + 1, 0)
        price_window = self.original_data[start_idx:self.data_index + 1, 3]  # 取close价

        if len(price_window) >= 2:
            high = price_window.max()
            low = price_window.min()
            if low > 0:
                volatility_ratio = (high - low) / low
                if volatility_ratio > volatility_threshold:
                    return True

        return False

    def step(self, action):
        current_price = self.original_data[self.data_index][3]  # 收盘价
        self.account.set_price(current_price)
        action_mask = self.get_action_mask()
        reward = 0.0
        terminated = False

        is_invalid_action = not action_mask[action]
        if is_invalid_action:
            self.bout_invalid_actions += 1

        open = False
        close = False
        if not is_invalid_action:
            if action == 1:
                open = self.account.open_long()
            elif action == 2:
                open = self.account.open_short()
            elif action == 3:
                close = self.account.close_long()
            elif action == 4:
                close = self.account.close_short()

        self.account.update_net_worth()
        if close:
            reward_value = close / self.account.initial_balance * 10_000
            reward += reward_value
            self._record_reward_event('平仓盈亏奖惩', reward_value)
        if open:
            reward_value = open / self.account.initial_balance * 10_000
            reward += reward_value
            self._record_reward_event('开仓手续费惩罚', reward_value)

        net_worth_change = (self.account.net_worth - self.account.old_net_worth) / self.account.initial_balance * 10_000
        if abs(net_worth_change) > 0:  # 只记录有实际变化的情况
            reward += net_worth_change
            self._record_reward_event('净值变动奖励', net_worth_change)

        # drawdown = self.account.get_drawdown()
        # if drawdown < 0:
        #     drawdown *= 10_000
        #     reward += drawdown
        #     self._record_reward_event('回撤处罚', drawdown)


        gain_ratio = self.account.get_gain_ratio()
        if gain_ratio >= 0.01:
            self._record_bout_event("止盈")
            terminated = True
        elif gain_ratio <= -0.005:
            self._record_bout_event("止损")
            terminated = True
        elif self.bout_total_steps >= 1000 and self.bout_invalid_actions > self.bout_total_steps * 0.10:
            self._record_bout_event("过度无效")
            terminated = True
        # elif self.account.order_nums > ((self.bout_total_steps // 86400) + 1) * 20:
        #     self._record_bout_event("过多下单终止")
        #     terminated = True
        # elif self.account.order_nums >= 20 and gain_ratio < 0.01:
        #     self._record_bout_event("未达到要求")
        #     terminated = True
        elif self.data_index + 1 >= len(self.observation_data) - self.seq_len - 1:
            self._record_bout_event("数据末尾")
            terminated = True


        self.account_state_history[self.data_index] = self.account.get_account_state() # 保存账户状态
        self.action_count[action] += 1 # 记录操作计数
        self.data_index += 1
        self.bout_total_steps += 1
        self.bout_total_reward += reward

        if terminated:
            self.render()

        infos = dict(
            is_invalid_action = is_invalid_action
        )

        return self.get_obs(), reward, terminated, False, infos

    def render(self):
        print(pretty_format_dict({
            '---------------------------------------': '',
            "步数": self.bout_total_steps,
            "无效": self.bout_invalid_actions,
            '总资产': float(self.account.net_worth),
            "总单数": self.account.order_nums,
            "盈利单数": self.account.win_order_nums,
            "亏损单数": self.account.loss_order_nums,
            '总分': float(self.bout_total_reward),
            '分数统计': self.reward_stats,  # 添加统计信息到info
            '操作统计': self.action_count,  # 输出操作统计
            '结束回合': self.bout_stats,
        }))
def pretty_format_dict(d, indent=0):
    lines = []
    for key, value in d.items():
        prefix = "  " * indent + f"{key}:"
        if isinstance(value, dict):
            lines.append(prefix)
            lines.append(pretty_format_dict(value, indent + 1))
        else:
            lines.append(f"{prefix} {value}")
    return "\n".join(lines)